package com.qust.web.controller;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.qust.web.dto.ResponseDto;
import com.qust.web.entity.ImgInfo4Redis;
import com.qust.web.vo.SplitResVo;
import lombok.SneakyThrows;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.multipart.MultipartFile;

import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.UUID;
import java.util.concurrent.*;

/**
 * @author Frank KONG
 * @version 1.0
 * @date 2021/11/26 16:02
 */
@RestController
@RequestMapping("/split")
public class SplitController {

    /**
     * 自定义线程池
     */
    private static ThreadPoolExecutor executor = new ThreadPoolExecutor(
            1,
            4,
            1L,
            TimeUnit.SECONDS,
            new ArrayBlockingQueue<>(1000)
    );

    @Value("${web.file.path}")
    private String imgWebPath;

    @Value("${redis.imgList.key}")
    private String imgList;

    @Autowired
    private RedisTemplate<String, String> redisTemplate;

    @RequestMapping(value = "/upload", method = RequestMethod.POST, produces = "application/json")
    public ResponseDto split(@RequestParam(value = "inputFile", required = true) MultipartFile inputFile,
                             @RequestParam(value = "models", required = true) String models) throws IOException, InterruptedException {

        //当前应用运行路径 classpath
        ClassLoader classLoader = this.getClass().getClassLoader();
        String path = classLoader.getResource("").getPath();

        //保存文件
        File targetFile = new File(path + "public/file/img/" + inputFile.getOriginalFilename());
        inputFile.transferTo(targetFile);

        //线程池执行任务
        List<SplitResVo> resVos = new ArrayList<>();
        String[] split = models.split(",");
        CountDownLatch latch = new CountDownLatch(split.length);
        for (int i = 0; i < split.length; i++) {
            executor.execute(new MyTask(split[i], split[i], imgWebPath + inputFile.getOriginalFilename(), resVos, latch));
        }

        //主线程等待任务结束
        latch.await(60, TimeUnit.SECONDS);
        return ResponseDto.ok(resVos);
    }

    /**
     * 自定义工作任务类
     */
    class MyTask implements Runnable {

        //要调用的模型名称
        String model;
        //list的key
        String key;
        //图片的路径
        String webPath;
        //并发计算后，结果存储的地方，需要加锁
        List<SplitResVo> resList;
        //CountDownLatch
        CountDownLatch latch;

        public MyTask(String model, String key, String webPath, List<SplitResVo> resList, CountDownLatch latch) {
            this.model = model;
            this.key = key;
            this.webPath = webPath;
            this.resList = resList;
            this.latch = latch;
        }

        @SneakyThrows
        @Override
        public void run() {
            try {
                //待处理图像加入到redis队列中
                String id = UUID.randomUUID().toString();
                ImgInfo4Redis imgInfo4Redis = new ImgInfo4Redis(id, webPath);
                ObjectMapper mapper = new ObjectMapper();
                String value = mapper.writeValueAsString(imgInfo4Redis);
                redisTemplate.opsForList().leftPush(key, value);
                //等待GPU服务器返回结果
                String res = null;
                long ttl = 0L;
                while (true) {
                    res = redisTemplate.opsForValue().get(id);
                    if(res == null){
                        ttl += 500;
                        Thread.sleep(500);
                    }else {
                        break;
                    }
                    if(ttl >= 10000) break;
                }
                if(res != null) {
                    synchronized (resList) {
                        resList.add(new SplitResVo(res, model));
                    }
                }
            }catch (Exception e){
                e.printStackTrace();
            }finally {
                latch.countDown();
            }
        }
    }

}
